import pandas as pd
import argparse
import os
import re

def clean_data(input_file, output_file, dup_cols=None, date_cols=None, lower_cols=None):
    """
    数据清洗主函数
    :param input_file: 输入文件
    :param output_file: 输出文件
    :param dup_cols: 用于判断重复的列
    :param date_cols: 需要标准化的日期列
    :param lower_cols: 需要转为小写的列
    """
    try:
        # 读取数据
        ext = os.path.splitext(input_file)[1].lower()
        if ext == '.csv':
            df = pd.read_csv(input_file)
        else:  # Excel文件
            df = pd.read_excel(input_file)
        
        original_rows = len(df)
        print(f"读取数据完成，共 {original_rows} 行")
        
        # 去除完全空行
        df.dropna(how='all', inplace=True)
        print(f"去除空行后: {len(df)} 行")
        
        # 去重
        if dup_cols:
            dup_columns = [col.strip() for col in dup_cols.split(',')]
            # 检查列是否存在
            valid_cols = [col for col in dup_columns if col in df.columns]
            if valid_cols:
                df.drop_duplicates(subset=valid_cols, keep='first', inplace=True)
                print(f"按 {valid_cols} 去重后: {len(df)} 行")
            else:
                print(f"指定的去重列不存在: {dup_columns}")
        
        # 日期标准化
        if date_cols:
            date_columns = [col.strip() for col in date_cols.split(',')]
            for col in date_columns:
                if col in df.columns:
                    try:
                        # 转换为日期格式
                        df[col] = pd.to_datetime(df[col], errors='coerce')
                        # 格式化为YYYY-MM-DD
                        df[col] = df[col].dt.strftime('%Y-%m-%d')
                        print(f"已标准化日期列: {col}")
                    except Exception as e:
                        print(f"标准化日期列 {col} 失败: {str(e)}")
        
        # 转为小写
        if lower_cols:
            lower_columns = [col.strip() for col in lower_cols.split(',')]
            for col in lower_columns:
                if col in df.columns and df[col].dtype == 'object':
                    df[col] = df[col].str.lower()
                    print(f"已转为小写列: {col}")
        
        # 去除字符串前后空格
        for col in df.select_dtypes(include=['object']).columns:
            df[col] = df[col].str.strip()
        
        # 保存清洗后的数据
        if output_file.lower().endswith('.csv'):
            df.to_csv(output_file, index=False, encoding='utf-8-sig')
        else:
            df.to_excel(output_file, index=False)
        
        print(f"\n数据清洗完成，已保存到: {output_file}")
        print(f"共处理 {original_rows - len(df)} 行异常数据")
        
    except Exception as e:
        print(f"数据清洗失败: {str(e)}")

if __name__ == "__main__":
    # 检查依赖
    try:
        import pandas
        import openpyxl
    except ImportError:
        print("请安装依赖: pip install pandas openpyxl")
        exit(1)
    
    parser = argparse.ArgumentParser(description='数据清洗工具')
    parser.add_argument('input', help='输入文件（Excel或CSV）')
    parser.add_argument('output', help='输出文件（Excel或CSV）')
    parser.add_argument('--dup-cols', help='用于去重的列，逗号分隔')
    parser.add_argument('--date-cols', help='需要标准化的日期列，逗号分隔')
    parser.add_argument('--lower-cols', help='需要转为小写的列，逗号分隔')
    
    args = parser.parse_args()
    
    clean_data(
        input_file=args.input,
        output_file=args.output,
        dup_cols=args.dup_cols,
        date_cols=args.date_cols,
        lower_cols=args.lower_cols
    )
